import numpy as np
from PyEMD import EMD, Visualisation
import matplotlib.pyplot as plt
# 构建信号
# 时间t: 为0到1s,采样频率为100Hz，S为合成信号
from pandas import read_csv, DataFrame
from sklearn.preprocessing import MinMaxScaler
from statsmodels.graphics.tsaplots import plot_pacf
from tensorflow.python.keras.utils.np_utils import to_categorical
from keras.layers import LSTM, RepeatVector, Dense, \
    Activation, Add, Reshape, Input, Lambda, Multiply, Concatenate, Dot
from util import series_to_supervised, TIME_STEP, draw_data
import keras
from sklearn.metrics import mean_squared_error, mean_absolute_error
from tensorflow.python.keras.callbacks import EarlyStopping
def getFlowData():
    # 切换19年和20年数据只需要改下面两行
    multi_dataset = read_csv('./data/2019allday.csv', header=0, index_col=0)
    day_num = 26  # 数据包含的天数
    # multi_dataset = read_csv('./data/多源数据总表.csv', header=0, index_col=None)
    # day_num = 31  # 数据包含的天数

    dataset = DataFrame()
    # 取in_card_flow(流出机场客流)、实际降落载客数 arr_ALDT_passenger、时段、天气作为参数，预测in_card_flow
    dataset['in_flow'] = multi_dataset['in_flow']
    # 将NAN替换为0
    dataset.fillna(0, inplace=True)
    values = dataset.values.flatten()
    return values

# PARAMETER_NUM: 分解出的本征函数个数+残波个数
def getModel():
    input1 = Input(shape=(TIME_STEP, 1))  # 输入时间序列数据
    lstm1 = keras.layers.Bidirectional(LSTM(4, return_sequences=False))(input1)
    input2 = Input(shape=(TIME_STEP, 1))  # 输入时间序列数据
    lstm2 = keras.layers.Bidirectional(LSTM(4, return_sequences=False))(input2)
    input3 = Input(shape=(TIME_STEP, 1))  # 输入时间序列数据
    lstm3 = keras.layers.Bidirectional(LSTM(4, return_sequences=False))(input3)
    input4 = Input(shape=(TIME_STEP, 1))  # 输入时间序列数据
    lstm4 = keras.layers.Bidirectional(LSTM(4, return_sequences=False))(input4)
    input5 = Input(shape=(TIME_STEP, 1))  # 输入时间序列数据
    lstm5 = keras.layers.Bidirectional(LSTM(4, return_sequences=False))(input5)
    input6 = Input(shape=(TIME_STEP, 1))  # 输入时间序列数据
    lstm6 = keras.layers.Bidirectional(LSTM(4, return_sequences=False))(input6)
    input7 = Input(shape=(TIME_STEP, 1))  # 输入时间序列数据
    lstm7 = keras.layers.Bidirectional(LSTM(4, return_sequences=False))(input7)
    input8 = Input(shape=(TIME_STEP, 1))  # 输入时间序列数据
    lstm8 = keras.layers.Bidirectional(LSTM(4, return_sequences=False))(input8)
    input9 = Input(shape=(TIME_STEP, 1))  # 输入时间序列数据
    lstm9 = keras.layers.Bidirectional(LSTM(4, return_sequences=False))(input9)
    input10 = Input(shape=(TIME_STEP, 1))  # 输入时间序列数据
    lstm10 = keras.layers.Bidirectional(LSTM(4, return_sequences=False))(input10)

    x = keras.layers.concatenate([lstm1, lstm2, lstm3, lstm4, lstm5, lstm6, lstm7, lstm8, lstm9, lstm10])

    output = Dense(1)(x)
    output = Reshape((-1,))(output)
    output = Dense(1)(output)  # 最后预测，代码修改点

    model = keras.models.Model(inputs=[input1, input2, input3, input4, input5, input6, input7, input8, input9, input10], outputs=output)
    return model


def get_trained_model():
    # 获取轨道客流数据
    flowData = getFlowData()
    train_flow_data = flowData.tolist()[:2500]
    # 提取imfs和残波
    emd = EMD()
    emd.emd(np.array(train_flow_data))
    imfs, res = emd.get_imfs_and_residue()
    # 绘制 IMF
    # t = np.arange(0, 2600, 1)
    # vis = Visualisation()
    # vis.plot_imfs(imfs=imfs, residue=res, t=t, include_residue=True)
    # # 绘制并显示所有提供的IMF的瞬时频率
    # vis.plot_instant_freq(t, imfs=imfs)
    # vis.show()
    # 偏自相关检验，确定lstm预测时的步长,暂时不知道怎么确定
    # for imf in imfs:
    #     plot_pacf(imf)
    # plt.show()
    # 把所有的本征函数和残波当做输入特征
    values = np.vstack((imfs, res))

    # 把输入整理成 n个time_step的形式，输出是下一个时间片的真实客流
    # 因为res残波全是0，所以这里没有用到 values[10]
    input1 = series_to_supervised(values[0].tolist(), TIME_STEP, 1)
    input2 = series_to_supervised(values[1].tolist(), TIME_STEP, 1)
    input3 = series_to_supervised(values[2].tolist(), TIME_STEP, 1)
    input4 = series_to_supervised(values[3].tolist(), TIME_STEP, 1)
    input5 = series_to_supervised(values[4].tolist(), TIME_STEP, 1)
    input6 = series_to_supervised(values[5].tolist(), TIME_STEP, 1)
    input7 = series_to_supervised(values[6].tolist(), TIME_STEP, 1)
    input8 = series_to_supervised(values[7].tolist(), TIME_STEP, 1)
    input9 = series_to_supervised(values[8].tolist(), TIME_STEP, 1)


    input1.drop(input1.columns[[-1]], axis=1, inplace=True)
    input2.drop(input2.columns[[-1]], axis=1, inplace=True)
    input3.drop(input3.columns[[-1]], axis=1, inplace=True)
    input4.drop(input4.columns[[-1]], axis=1, inplace=True)
    input5.drop(input5.columns[[-1]], axis=1, inplace=True)
    input6.drop(input6.columns[[-1]], axis=1, inplace=True)
    input7.drop(input7.columns[[-1]], axis=1, inplace=True)
    input8.drop(input8.columns[[-1]], axis=1, inplace=True)
    input9.drop(input9.columns[[-1]], axis=1, inplace=True)


    # 构建输出数据
    output_train = train_flow_data[TIME_STEP:]
    output_train = np.array(output_train)

    model = getModel()
    model.compile(loss='mse', optimizer='adam')
    model.fit(
        [input1, input2, input3, input4, input5, input6, input7, input8,
         input9], output_train, verbose=2,
        epochs=200, batch_size=128)
    return model

if __name__ == '__main__':
    # 获取轨道客流数据
    flowData = getFlowData()
    test_flow_data = flowData.tolist()[2500:]
    # 提取imfs和残波
    emd = EMD()
    emd.emd(np.array(test_flow_data))
    imfs, res = emd.get_imfs_and_residue()
    # 把所有的本征函数和残波当做输入特征
    values = np.vstack((imfs, res))


    # 把输入整理成 n个time_step的形式，输出是下一个时间片的真实客流
    # 因为res残波全是0，所以这里没有用到 values[10]
    input1 = series_to_supervised(values[0].tolist(), TIME_STEP, 1)
    input2 = series_to_supervised(values[1].tolist(), TIME_STEP, 1)
    input3 = series_to_supervised(values[2].tolist(), TIME_STEP, 1)
    input4 = series_to_supervised(values[3].tolist(), TIME_STEP, 1)

    input1.drop(input1.columns[[-1]], axis=1, inplace=True)
    input2.drop(input2.columns[[-1]], axis=1, inplace=True)
    input3.drop(input3.columns[[-1]], axis=1, inplace=True)
    input4.drop(input4.columns[[-1]], axis=1, inplace=True)


    # 构建输出数据
    output_test = test_flow_data[TIME_STEP:]



    # 重塑成3D形状 [样例, 时间步, 特征]
    input1_test = input1.values.reshape((input1.shape[0], TIME_STEP, int(input1.shape[1] / TIME_STEP)))
    input2_test = input2.values.reshape((input2.shape[0], TIME_STEP, int(input2.shape[1] / TIME_STEP)))
    input3_test = input3.values.reshape((input3.shape[0], TIME_STEP, int(input3.shape[1] / TIME_STEP)))
    input4_test = input4.values.reshape((input4.shape[0], TIME_STEP, int(input4.shape[1] / TIME_STEP)))




    # model.fit函数中传入的输入输出，必须是ndarray格式
    output_test = np.array(output_test)


    model = get_trained_model()
    yhat = model.predict([input1_test, input2_test, input3_test, input4_test, input5_test, input6_test, input7_test, input8_test, input9_test, input10_test])
    mse = mean_squared_error(output_test, yhat)
    print('Test MSE: %.3f' % mse)
    pic_path = "./picture/EMD+LSTM.png"
    pic_title = "EMD\n MSE=%.3f" % mse
    draw_data(yhat, output_test, pic_path, pic_title)
